package bitmap

import (
	"strings"
	"sync"
)

type SafeLenBitmap struct {
	bits []byte
	len  int
	mux  sync.Mutex
}

// 初始化一个位图
func NewSafeLenBitmap(max int) *SafeLenBitmap {
	size := max >> shift
	if max&mask > 0 {
		size += 1
	}

	return &SafeLenBitmap{
		bits: make([]byte, size),
	}
}

// 用二进制初始化一个位图
func NewSafeLenBitmapWithBinary(b []byte) *SafeLenBitmap {
	bt := &SafeLenBitmap{
		bits: make([]byte, len(b)),
	}
	copy(bt.bits, b)

	for i := 0; i < len(bt.bits); i++ {
		for j := uint(0); j < bitSize; j++ {
			if bt.bits[i]&(1<<j) == 1 {
				bt.len++
			}
		}
	}

	return bt
}

// 用二进制字符串初始化一个位图
func NewSafeLenBitmapWithBinaryString(s string) *SafeLenBitmap {
	l := len(s)
	size := l >> shift
	if l&mask > 0 {
		size += 1
	}

	bt := &SafeLenBitmap{
		bits: make([]byte, size),
	}

	for i := 0; i < len(s); i++ {
		if s[i] == '0' {
			bt.Unset(i)
			continue
		}
		if s[i] == '1' {
			bt.Set(i)
			continue
		}
	}

	return bt
}

func (b *SafeLenBitmap) Len() int {
	return b.len
}

func (b *SafeLenBitmap) Cap() int {
	return cap(b.bits)
}

// 添加数字
func (b *SafeLenBitmap) Set(num int) *SafeLenBitmap {
	b.setBit(num, true)
	return b
}

// 清除数字
func (b *SafeLenBitmap) Unset(num int) *SafeLenBitmap {
	b.setBit(num, false)
	return b
}

// 数字是否在位图中
func (b *SafeLenBitmap) Exist(num int) bool {
	index, pos := indexAndPos(num)
	return index < len(b.bits) && (b.bits[index]&(1<<uint(pos)) != 0)
}

// 设置bit
func (b *SafeLenBitmap) setBit(num int, stat bool) *SafeLenBitmap {
	b.mux.Lock()
	defer b.mux.Unlock()

	index, pos := indexAndPos(num)
	b.upCapToIndex(index)
	if stat {
		b.bits[index] |= 1 << uint(pos)
		b.len++
	} else {
		b.bits[index] &^= 1 << uint(pos)
		b.len--
	}

	return b
}

// 扩容
func (b *SafeLenBitmap) upCapToIndex(index int) {
	for index+1 > len(b.bits) {
		b.bits = append(b.bits, 0)
	}
}

// 2倍扩容
// func (b *SafeLenBitmap) upCapToIndex(index int) {
// 	toIndex := index + 1
// 	l := len(b.bits)
// 	if l < toIndex {
// 		l2 := 2 * l
// 		if toIndex < l2 {
// 			toIndex = l2
// 		}
// 		n := make([]byte, toIndex, toIndex)
// 		copy(n, b.bits)
// 		b.bits = n
// 	}
// }

// 二进制
func (b *SafeLenBitmap) Byte() []byte {
	return b.bits
}

// 二进制字符串
func (b *SafeLenBitmap) String() string {
	var sb strings.Builder
	for i := 0; i < len(b.bits); i++ {
		sb.WriteString(byteToString(b.bits[i]))
	}

	return sb.String()
}
